#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue


template<typename TYPE_X, typename TYPE_Z> 
class KernelIm2Col {
    using T = TYPE_X;
public:
    __aicore__ inline KernelIm2Col() {}

    __aicore__ inline uint32_t AlignUp(uint32_t a, uint32_t b) 
    {
        if (b == 0)
            return a;
        return ((a + 32 / b - 1) / (32 / b)) * (32 / b);
    }
    
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR z, const Im2ColCustomTilingData &tiling_data)
    {
        this->tiling = tiling_data;

        ASSERT(AscendC::GetBlockNum() != 0 && "block dim can not be zero!");
        this->core_id = AscendC::GetBlockIdx();
        // 如果当前核没有分配到向量，直接返回
        if (tiling.core_element_count[this->core_id] <= 0) {
            return ;
        }

        // 输入和输出GM地址设置
        xGm.SetGlobalBuffer((__gm__ TYPE_X *)x, tiling.input_elements);
        zGm.SetGlobalBuffer((__gm__ TYPE_Z *)z + tiling.core_element_start[this->core_id], tiling.core_element_count[this->core_id]);
        
        // 只需要输出缓冲区，直接从GM读取输入
        this->aligned_tile_size = AlignUp((uint32_t)tiling.tile_element_num, sizeof(TYPE_Z));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->aligned_tile_size * sizeof(TYPE_X));
    }
    
    __aicore__ inline void Process()
    {
        int32_t loopCount = tiling.core_loop_times[this->core_id];
        
        // 处理完整的循环
        for (int32_t i = 0; i < loopCount; i++) 
        {
            // 判断是否是最后一次循环
            if (i == loopCount - 1) {
                // 最后一次循环，处理tail
                this->processDataNum = AlignUp((uint32_t)tiling.core_tail_elements[this->core_id], sizeof(TYPE_Z));
            } else {
                // 正常循环
                this->processDataNum = this->aligned_tile_size;
            }
            
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<TYPE_Z> zLocal = outQueueZ.AllocTensor<TYPE_Z>();

        // 计算当前tile在输出矩阵中的起始位置
        int32_t current_tile_size = (progress == tiling.core_loop_times[this->core_id] - 1) ? 
                                    tiling.core_tail_elements[this->core_id] : 
                                    tiling.tile_element_num;
        int32_t outputStart = tiling.core_element_start[this->core_id] + progress * tiling.tile_element_num;
        
        // 为当前tile生成im2col数据
        for (int32_t i = 0; i < current_tile_size; i++) {
            int32_t outputIdx = outputStart + i;
            
            // 边界检查
            if (outputIdx >= tiling.total_output_elements) {
                zLocal.SetValue(i, 0.0f);
                continue;
            }
            
            // 计算输出索引对应的输入位置
            int32_t inputIdx = CalculateInputIndex(outputIdx);
            
            if (inputIdx >= 0 && inputIdx < tiling.input_elements) {
                // 直接从GM读取
                TYPE_X value = xGm.GetValue(inputIdx);
                zLocal.SetValue(i, value);
            } else {
                // 填充区域设为0
                zLocal.SetValue(i, 0.0f);
            }
        }
        
        outQueueZ.EnQue<TYPE_Z>(zLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<TYPE_Z> zLocal = outQueueZ.DeQue<TYPE_X>();
        
        int32_t current_tile_size = (progress == tiling.core_loop_times[this->core_id] - 1) ? 
                                    tiling.core_tail_elements[this->core_id] : 
                                    tiling.tile_element_num;
        int32_t outputOffset = progress * tiling.tile_element_num;
        
        // 边界检查
        int32_t remaining = tiling.core_element_count[this->core_id] - outputOffset;
        if (current_tile_size > remaining) {
            current_tile_size = remaining;
        }
        
        uint32_t bytes_to_copy = current_tile_size * sizeof(TYPE_Z);
        AscendC::DataCopyExtParams copyParams{
                    1,                          // blockCount
                    bytes_to_copy,              // blockLen (字节数)
                    0,                          // srcStride
                    0,                          // dstStride
                    0                           // rsv
                };
                
        AscendC::DataCopyPad(zGm[outputOffset], zLocal, copyParams);
        
        outQueueZ.FreeTensor(zLocal);
    }

    // 计算输出索引对应的输入索引
    __aicore__ inline int32_t CalculateInputIndex(int32_t outputIdx)
    {
        // 将输出展平索引转换为 [N, output_channels, L] 坐标
        int32_t n_idx = outputIdx / (tiling.output_channels * tiling.L);
        int32_t remaining = outputIdx % (tiling.output_channels * tiling.L);
        int32_t channel_block = remaining / tiling.L;  // 0 ~ (C*kH*kW-1)
        int32_t pos_idx = remaining % tiling.L;        // 0 ~ (L-1)
        
        // 从channel_block中分解出c, kh, kw
        int32_t c_idx = channel_block / (tiling.kernel_h * tiling.kernel_w);
        int32_t kernel_idx = channel_block % (tiling.kernel_h * tiling.kernel_w);
        int32_t kh_idx = kernel_idx / tiling.kernel_w;
        int32_t kw_idx = kernel_idx % tiling.kernel_w;
        
        // 从pos_idx中分解出out_h, out_w
        int32_t out_h = pos_idx / tiling.out_W;
        int32_t out_w = pos_idx % tiling.out_W;
        
        // 计算对应的输入位置
        int32_t in_h = out_h * tiling.stride_val + kh_idx - tiling.padding_val;
        int32_t in_w = out_w * tiling.stride_val + kw_idx - tiling.padding_val;
        
        // 检查边界
        if (in_h < 0 || in_h >= tiling.H || in_w < 0 || in_w >= tiling.W) {
            return -1;  // 表示填充位置
        }
        
        // 计算输入张量中的索引
        return n_idx * (tiling.C * tiling.H * tiling.W) + c_idx * (tiling.H * tiling.W) + in_h * tiling.W + in_w;
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::GlobalTensor<TYPE_X> xGm;
    AscendC::GlobalTensor<TYPE_X> zGm;

    int32_t core_id;
    int32_t aligned_tile_size;
    int32_t processDataNum;

    Im2ColCustomTilingData tiling;
};

extern "C" __global__ __aicore__ void im2_col_custom(GM_ADDR x, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling); 
    KernelIm2Col<DTYPE_X, DTYPE_Z> op;
    op.Init(x, z, tiling_data);
    op.Process();
}

